package middleware

import (
	"bytes"
	"context"
	"errors"
	"io"
	"net/http"
	"paper-pro/pkg/result"
	"paper-pro/pkg/xerr"
	"path/filepath"
	"strings"

	"github.com/zeromicro/go-zero/rest/httpx"
)

var (
	ErrCheckFileTypeFailed = xerr.NewErrCodeMsg(xerr.REQUEST_PARAM_ERROR, "only PDF files are allowed")
	ErrCheckFileSizeFailed = xerr.NewErrCodeMsg(xerr.REQUEST_PARAM_ERROR, "file too large (max 20MB)")
	ErrInvalidFileName     = xerr.NewErrCodeMsg(xerr.REQUEST_PARAM_ERROR, "invalid file name (cannot contain .. or / or \\)")
)

var pdfMagicNumber = []byte("%PDF-") // 新增PDF魔数验证

// 自定义ReadCloser，确保关闭时释放原始part资源
type combinedReadCloser struct {
	io.Reader
	closer io.Closer
}

func (c *combinedReadCloser) Close() error {
	return c.closer.Close() // 调用原始part的Close方法
}

type FileCheckMiddleware struct {
	maxFileSize int64
}

func NewFileCheckMiddleware(maxFileSize int64) *FileCheckMiddleware {
	return &FileCheckMiddleware{
		maxFileSize: maxFileSize,
	}
}

func (m *FileCheckMiddleware) Handle(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		// 1. 验证Content-Type
		if !strings.Contains(r.Header.Get("Content-Type"), "multipart/form-data") {
			httpx.Error(w, errors.New("invalid content type"))
			return
		}

		// 2. 验证文件大小。通过Content-Length预判（注意：需配合WithMaxBytes使用）
		if r.ContentLength > m.maxFileSize {
			result.HttpResult(r, w, nil, ErrCheckFileSizeFailed)
			return
		}

		// 3. 使用流式处理替代ParseMultipartForm（内存优化）
		reader, err := r.MultipartReader()
		if err != nil {
			result.HttpResult(r, w, nil, err)
			return
		}

		// 4. 增强文件类型验证（魔数+Content-Type双验证）和文件名
		part, err := reader.NextPart()
		if err != nil || part.FormName() != "file" {
			result.HttpResult(r, w, nil, ErrCheckFileTypeFailed)
			return
		}

		// 验证文件名安全性
		filename := part.FileName()
		if strings.Contains(filename, "..") || strings.Contains(filename, "/") || strings.Contains(filename, "\\") {
			result.HttpResult(r, w, nil, ErrInvalidFileName)
			return
		}

		// 读取文件前512字节用于魔数验证
		buf := make([]byte, 512)
		n, err := part.Read(buf)
		if err != nil && err != io.EOF {
			result.HttpResult(r, w, nil, err)
			return
		}

		// 验证PDF魔数和Content-Type
		if !bytes.HasPrefix(buf[:n], pdfMagicNumber) ||
			!strings.Contains(part.Header.Get("Content-Type"), "application/pdf") {
			result.HttpResult(r, w, nil, ErrCheckFileTypeFailed)
			return
		}

		// 创建新的Reader组合已读数据和剩余内容
		combinedReader := io.MultiReader(
			bytes.NewReader(buf[:n]),
			part,
		)

		// 替换请求体供后续处理器使用
		r.Body = &combinedReadCloser{
			Reader: combinedReader,
			closer: part, // 传递原始part作为closer
		}

		// 在验证通过后添加文件信息
		ctx := r.Context()
		ctx = context.WithValue(ctx, "fileMeta", map[string]interface{}{
			"filename": filepath.Base(filename), // 使用path.Base确保安全文件名
			"size":     r.ContentLength,
			"type":     part.Header.Get("Content-Type"),
		})
		r = r.WithContext(ctx)

		// Passthrough to next handler if need
		next(w, r)
	}
}
